import sys
from functools import partial

from PyQt6 import uic
from PyQt6.QtCore import QModelIndex, Qt
from PyQt6.QtGui import QIcon, QStandardItemModel, QStandardItem
from PyQt6.QtWidgets import QWidget, QTreeView, QAbstractItemView, QTableWidget, QTableWidgetItem, QCheckBox, \
    QPushButton
from EXFAT.Constant import *
from EXFAT.EXFAT import EXFAT
from pojo.Disk import Disk
from pojo.File import File
from EXFAT.Constant import SUFFIX_LIST
from Utils.CommonUtils import open_with_explorer
from MessageBox.PathWriteErrorMessageBox import PathWriteErrorMessageBox
from MessageBox.ProcessingBox import ProcessingBox


class ScanFileWindow:
    def __init__(self, disk: Disk, parent):
        self.parent = parent
        self.tree = None
        self.disk = disk
        self.cache_log = []
        self.root_log = []
        self.current_log: list[File] = []
        self.checked_data = {}

        self.ui: QWidget = uic.loadUi("./UI/scan.ui")
        self.ui.setWindowIcon(QIcon("./FormatRecovery.ico"))
        self.ui.setWindowTitle("扫描结果")
        self.ui.show()

        self.tree: QTreeView = self.ui.findChild(QTreeView, "treeView")
        self.v_model = QStandardItemModel()
        self.tree.setModel(self.v_model)
        self.tree.header().setVisible(False)
        self.model = QStandardItem()
        self.model.setText(self.parent.current_disk)
        self.v_model.appendRow(self.model)
        self.tree.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        self.tree.clicked.connect(self.treeViewClicked)
        self.v_model.dataChanged.connect(self.on_data_changed)
        self.v_model.itemChanged.connect(self.folder_checked)

        self.table: QTableWidget = self.ui.findChild(QTableWidget, "tableWidget")
        self.table.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        self.table.setColumnWidth(0, 30)
        self.table.setColumnWidth(2, 80)
        self.table.horizontalHeader().sectionClicked.connect(self.all_selected)

        self.initButton()
        self.initTreeView(self.disk.file_log, self.model)

    def cancel(self):
        sys.exit()

    def next_step(self):
        self.disk: EXFAT
        self.disk.back_to_cluster_table()
        path: str = open_with_explorer()
        if path is not None and len(self.checked_data) > 0:
            process = ProcessingBox()
            values: list[File] = list(self.checked_data.values())
            values.sort(key=lambda i: i.local)
            process.set_offset(int(100 / len(values)))

            def work():
                self.disk.write_to_path(values, path, process)
                process.callback()
            work()

            if process.res == SAVE_SUCCESS:
                # TODO 弹出保存成功
                pass
            else:
                print("程序出错")
        else:
            PathWriteErrorMessageBox().show()

    def back(self):
        self.ui.close()
        self.parent.ui.show()

    def initButton(self):
        back: QPushButton = self.ui.findChild(QPushButton, "back")
        cancel: QPushButton = self.ui.findChild(QPushButton, "cancel")
        next_step: QPushButton = self.ui.findChild(QPushButton, "next")

        back.clicked.connect(lambda: self.back())
        cancel.clicked.connect(lambda: self.cancel())
        next_step.clicked.connect(lambda: self.next_step())

    def up_to_cancel_selected(self, parent: QStandardItem, item: QStandardItem):
        if parent is not None and parent.checkState() == Qt.CheckState.Checked:
            self.pop_checked_data_and_rebuild(parent.data(), item.data())
            self.v_model.blockSignals(True)
            parent.setCheckState(Qt.CheckState.Unchecked)
            self.v_model.blockSignals(False)

            self.cancel_data(parent.data())

            self.up_to_cancel_selected(parent.parent(), parent)

    def down_to_cancel_selected(self, item: QStandardItem):
        if item.hasChildren():
            for row in range(item.rowCount()):
                child: QStandardItem = item.child(row)
                self.v_model.blockSignals(True)
                child.setCheckState(Qt.CheckState.Unchecked)
                self.v_model.blockSignals(False)

                self.cancel_data(child.data())

                self.down_to_cancel_selected(child)

    def cancel_data(self, f: File):
        for i in f.getFiles():
            i.checked = False

    def to_cancel_selected(self, item: QStandardItem):
        parent: QStandardItem = item.parent()
        self.pop_checked_data(item.data())
        self.cancel_data(item.data())
        self.up_to_cancel_selected(parent, item)
        self.down_to_cancel_selected(item)

    def pop_checked_data_and_rebuild(self, f: File, excluded: File):
        if self.checked_data.get(f.name) is not None:
            data: File = self.checked_data.pop(f.name)
            if excluded.file_type == IS_FOLDER:
                iterator = data.getFolders()
            else:
                iterator = data.getFiles()
            for i in iterator:
                if i.name != excluded.name:
                    self.checked_data[i.name] = i

    def pop_checked_data(self, f: File):
        if f.name in list(self.checked_data.keys()):
            self.checked_data.pop(f.name)

    def on_data_changed(self):
        """手动触发页面更新"""
        self.tree.viewport().update()

    def set_status_all_child_items(self, parent_item: QStandardItem):
        """获取指定 QStandardItem 的所有子 item"""
        file: File = parent_item.data()
        self.pop_checked_data(file)
        if len(file.getFiles()) > 0:
            for i in file.getFiles():
                self.pop_checked_data(i)
                i.checked = True
        for row in range(parent_item.rowCount()):
            child_item = parent_item.child(row)
            self.v_model.blockSignals(True)
            child_item.setCheckState(Qt.CheckState.Checked)
            self.v_model.blockSignals(False)
            self.set_status_all_child_items(child_item)

    def folder_checked(self, item: QStandardItem):
        data: list or File = item.data()
        # 清理
        if isinstance(data, File):
            if item.checkState() == Qt.CheckState.Checked:
                self.set_status_all_child_items(item)
                self.checked_data[data.name] = data
            else:
                self.to_cancel_selected(item)

    def all_selected(self, index: int):
        parent = str(self.current_log[0].parent)
        checked = self.checked_data.get(parent) is None
        for i in range(self.table.rowCount()):
            checkbox: QCheckBox = self.table.cellWidget(i, 0)
            checkbox.blockSignals(True)
            checkbox.setChecked(checked)
            checkbox.blockSignals(False)
        if index == 0:
            for i in self.current_log:
                i: File
                if i.name in list[self.checked_data.keys()]:
                    self.checked_data.pop(i.name)
                i.checked = checked
            if self.checked_data.get(parent) is None:
                self.checked_data[parent] = self.current_log
            else:
                self.checked_data.pop(parent)

    def oneBoxCheck(self, state, file: File):
        if state == 2:
            self.checked_data[file.name] = file
            file.checked = True
        if state == 0:
            self.checked_data.pop(file.name)
            file.checked = False

    def refreshTableView(self, files: list[File]):
        if len(files) > 0:
            self.current_log = [i for i in files]
            self.table.setRowCount(len(files))
            for i, file in enumerate(files):
                checkbox = QCheckBox()
                if file.checked:
                    checkbox.setChecked(True)
                checkbox.stateChanged.connect(partial(self.oneBoxCheck, file=file))
                self.table.setCellWidget(i, 0, checkbox)
                self.table.setItem(i, 1, QTableWidgetItem(file.name))
                if file.file_type > 0:
                    file_type = SUFFIX_LIST[file.file_type]
                else:
                    file_type = "未知扩展名"
                self.table.setItem(i, 2, QTableWidgetItem(file_type))
                self.table.setItem(i, 3, QTableWidgetItem(str(file.size)))
                self.table.setItem(i, 4, QTableWidgetItem(""))
        self.cache_log.clear()

    def clearRows(self):
        while self.table.rowCount() > 0:
            self.table.removeRow(0)

    def headerClicked(self):
        self.refreshTableView(self.root_log)

    def treeViewClicked(self, index: QModelIndex):
        self.clearRows()
        item_data: list[File] or File = self.tree.model().itemFromIndex(index).data()
        if isinstance(item_data, list):
            for i in item_data:
                if i.file_type >= 0:
                    self.cache_log.append(i)
        else:
            item_data: File
            for i in item_data.getFiles():
                self.cache_log.append(i)
        self.refreshTableView(self.cache_log)

    def initTreeView(self, data: list[File], model: QStandardItem or QStandardItemModel, init=None):
        for i in data:
            i: File
            if i.hasChildren() and i.file_type == IS_FOLDER:
                item: QStandardItem = QStandardItem()
                item.setCheckable(True)
                item.setText(i.name)
                item.setData(i)
                model.appendRow(item)
                self.initTreeView(i.getFolders(), item, False)
            if i.file_type >= 0 and init is None:
                self.root_log.append(i)
        self.model.setData(self.root_log)
